import logging
import os
import time
import torch
import torch.nn as nn
from einops import rearrange

from tools.meter import AverageMeter
from torch.cuda import amp

def do_train(cfg,
             model,
             train_loader,
             val_loader,
             optimizer,
             scheduler,
             loss_fn, local_rank, saver):
    log_period = cfg.SOLVER.LOG_PERIOD
    checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
    eval_period = cfg.SOLVER.EVAL_PERIOD

    device = "cuda"
    epochs = cfg.SOLVER.MAX_EPOCHS

    logger = logging.getLogger("pit.train")
    logger.info('start training')
    _LOCAL_PROCESS_GROUP = None
    if device:
        model.to(local_rank)
        if torch.cuda.device_count() > 1 and cfg.MODEL.DIST_TRAIN:
            print('Using {} GPUs for training'.format(torch.cuda.device_count()))
            model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank],
                                                              find_unused_parameters=True)

    cls_loss_meter = AverageMeter()
    acc_meter = AverageMeter()

    scaler = amp.GradScaler()
    for epoch in range(1, epochs + 1):
        start_time = time.time()
        cls_loss_meter.reset()
        acc_meter.reset()
        scheduler.step(epoch)
        model.train()

        for n_iter, (img, vid) in enumerate(train_loader):
            optimizer.zero_grad()
            img = img.to(device)
            target = vid.to(device)
            with amp.autocast(enabled=True):
                img = rearrange(img, "B T C H W -> B C T H W")
                score = model(img)
                ID_LOSS = loss_fn(score, target)
                loss = cfg.MODEL.ID_LOSS_WEIGHT * ID_LOSS

            scaler.scale(loss).backward()

            scaler.step(optimizer)
            scaler.update()

            if isinstance(score, list):
                acc = (score[0][0].max(1)[1] == target).float().mean()
            else:
                acc = (score.max(1)[1] == target).float().mean()

            cls_loss_meter.update(ID_LOSS.item(), img.shape[0])
            acc_meter.update(acc, 1)
            torch.cuda.synchronize()
            if (n_iter + 1) % log_period == 0:
                logger.info(
                    "Epoch[{}] Iteration[{}/{}] cls_loss: {:.3f},  Acc: {:.3f}, Base Lr: {:.2e}"
                        .format(epoch, (n_iter + 1), len(train_loader),
                                cls_loss_meter.avg, acc_meter.avg, scheduler._get_lr(epoch)[0]))

        end_time = time.time()
        time_per_batch = (end_time - start_time) / (n_iter + 1)

        saver.dump_metric_tb(cls_loss_meter.avg, epoch, f'losses', f'cls_loss')
        saver.dump_metric_tb(1, epoch, f'losses', f'tri_loss')
        saver.dump_metric_tb(acc_meter.avg, epoch, f'losses', f'acc')
        saver.dump_metric_tb(optimizer.param_groups[0]['lr'], epoch, f'losses', f'lr')

        num_samples = cfg.SOLVER.IMS_PER_BATCH
        logger.info("Epoch {} done. Time per batch: {:.3f}[s] Speed: {:.1f}[samples/s]"
                    .format(epoch, time_per_batch, num_samples / time_per_batch))

        if epoch % checkpoint_period == 0:
            torch.save(model.state_dict(),
                       os.path.join(cfg.OUTPUT_DIR, "checkpoint", cfg.MODEL.NAME + '_{}.pth'.format(epoch)))

        val_flag = False

        if epoch % eval_period == 0:
            val_flag = True

        if not val_flag:
            continue

        total = 0
        top1 = 0
        top5 = 0
        top10 = 0

        acc_test_meter = AverageMeter()
        model.eval()
        for n_iter, (img, target) in enumerate(val_loader):
            with torch.no_grad():
                img = img.to(device)
                target = target.to(device)
                img = rearrange(img, "B T C H W -> B C T H W")
                score = model(img)

                # acc
                if isinstance(score, list):
                    acc = (score[0][0].max(1)[1] == target).float().mean()
                else:
                    acc = (score.max(1)[1] == target).float().mean()
                acc_test_meter.update(acc, 1)

                # top k
                _, maxk = torch.topk(score, 10, dim=-1)

                total += target.size(0)
                target = target.view(-1, 1)  # reshape labels from [n] to [n,1] to compare [n,k]

                top1 += (target == maxk[:, 0:1]).sum().item()
                top5 += (target == maxk[:, 0:5]).sum().item()
                top10 += (target == maxk[:, 0:10]).sum().item()


        logger.info("Validation Results acc {:.3f} - Epoch: {}".format(acc_meter.avg,epoch))
        logger.info("Top-{:<3}:{:.1%}".format(1, top1 / total))
        logger.info("Top-{:<3}:{:.1%}".format(5, top5 / total))
        logger.info("Top-{:<3}:{:.1%}".format(10, top10 / total))

    return None


def do_test(cfg,
            model,
            val_loader):
    device = "cuda"
    logger = logging.getLogger("pit.test")
    logger.info("Enter inferencing")

    if device:
        if torch.cuda.device_count() > 1:
            print('Using {} GPUs for inference'.format(torch.cuda.device_count()))
            model = nn.DataParallel(model)
        model.to(device)

    acc_meter = AverageMeter()
    model.eval()
    total = 0
    top1 = 0
    top5 = 0
    top10 = 0
    for n_iter, (img, target) in enumerate(val_loader):
        with torch.no_grad():
            img = img.to(device)
            img = rearrange(img, "B T C H W -> B C T H W")
            target = target.to(device)
            score = model(img)
            # acc
            if isinstance(score, list):
                acc = (score[0][0].max(1)[1] == target).float().mean()
            else:
                acc = (score.max(1)[1] == target).float().mean()
            acc_meter.update(acc, 1)

            # top k
            _, maxk = torch.topk(score, 10, dim=-1)

            total += target.size(0)
            target = target.view(-1, 1)  # reshape labels from [n] to [n,1] to compare [n,k]

            top1 += (target == maxk[:, 0:1]).sum().item()
            top5 += (target == maxk[:, 0:5]).sum().item()
            top10 += (target == maxk[:, 0:10]).sum().item()

    logger.info("Validation Results acc {:.3f} ".format(acc_meter.avg))
    logger.info("Top-{:<3}:{:.1%}".format(1, top1 / total))
    logger.info("Top-{:<3}:{:.1%}".format(5, top5 / total))
    logger.info("Top-{:<3}:{:.1%}".format(10, top10 / total))

